Skip to content

Conversation

dferre97
Copy link
Contributor

This PR adds support for efficient CUDA message passing on coalesced :coo graphs by exploiting SpMM. Performance is comparable to the case of :sparse graphs.

The changes introduced by this PR include:

  • Introduce a private __adjacency_matrix function, used during message passing, that efficiently constructs a CuSparseMatrixCOO for a :coo graph. This is leveraged by propagate to implement message passing via SpMM.
  • Implement a workaround for CUDA.jl issue #2820, which highlights that dense × sparse multiplication with a CuSparseMatrixCOO is computed incorrectly. The fix involves changing coalesce to sort the :coo graph representation by target instead of source.
  • Extend tests in GraphNeuralNetworks and GNNlib to cover coalesced :coo graphs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant